#include "vmlinux.h"
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>
#include <bpf/bpf_core_read.h>

#define MAX_PATH 4096
#define MAX_COMM 4096
#define NAME_MAX 256
#define MAX_PATH_DEPTH 16
#define RINGBUF_SIZE 1 << 22

#define ATTR_MODE	(1 << 0)
#define ATTR_UID	(1 << 1)
#define ATTR_GID	(1 << 2)

char _license[] SEC("license") = "GPL";

enum event_type {
    FILE_CREATE,
    FILE_DELETE,
    FILE_RENAME,
    FILE_CONTENT_CHANGED,
    FILE_PERM_CHANGE,
    MOUNT_ADDED,
    MOUNT_REMOVED
};

struct permission_data {
    int is_permission_changed;
};
struct path_info {
    int depth; // depth of the next path component
    int len; // Length of the path component
    u64 dev; // Device ID
    u64 ino; // Inode number
    char name[NAME_MAX];
    struct dentry *entry;
};

struct event {
    u32 pid;
    char comm[MAX_COMM];
    u32 type;
    u64 dev;
    struct path_info src_path[MAX_PATH_DEPTH];
    u64 src_ino;
    struct path_info dst_path[MAX_PATH_DEPTH];
    u64 dst_ino;
};

struct {
    __uint(type, BPF_MAP_TYPE_RINGBUF);
    __uint(max_entries, RINGBUF_SIZE);
} events SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_PERCPU_HASH);
    __uint(max_entries, MAX_PATH_DEPTH * 2);
    __type(key, u64);
    __type(value, struct path_info);
} tmp_pathinfos SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_PERCPU_HASH);
    __uint(max_entries, 1);
    __type(key, u64);
    __type(value, struct path_info);
} tmp_mount_pathinfos SEC(".maps");
struct {
    __uint(type, BPF_MAP_TYPE_PERCPU_HASH);
    __uint(max_entries, 1);
    __type(key, u64);
    __type(value, u64);
} permission_cache SEC(".maps");

/**
 * @brief 缓存路径信息
 *
 * 遍历dentry树以查找路径组件，并将路径信息缓存在tmp_pathinfos映射中。
 *
 * @param dentry dentry结构体指针
 * @param ino 文件的inode号
 * @param is_src 是否为源路径
 *
 * @return 缓存路径信息的结果，成功返回0，失败返回-1
 */
static int cache_path_infos(struct dentry *dentry, long unsigned int ino, bool is_src) {
    struct dentry *tmp_dentry = dentry;
    int depth = 0;
    // 遍历 dentry 树以查找路径组件
    // Traverse the dentry tree to find the path components
    for (int i = MAX_PATH_DEPTH - 1; i >= 0; i--) {
        struct path_info path_info;        
        path_info.dev = BPF_CORE_READ(dentry, d_sb, s_dev);
        path_info.depth = depth;
        path_info.entry = dentry;

        // 将 dentry 的名称复制到 path_info 结构中
        // Copy the name of the dentry to the path_info structure
        const unsigned char *name = BPF_CORE_READ(tmp_dentry, d_name.name);
        if (name == NULL) {
            return -1;
        }
        path_info.len = bpf_probe_read_kernel_str(path_info.name, NAME_MAX, name);
        // bpf_printk("path_info.name: %s\n", path_info.name);

        if (tmp_dentry == BPF_CORE_READ(tmp_dentry, d_parent)) {
            u64 index = is_src ? 0 : MAX_PATH_DEPTH;
            if (ino != 0) {
                path_info.ino = ino;
            } else {
                path_info.ino = BPF_CORE_READ(tmp_dentry, d_inode, i_ino);
            }
            bpf_map_update_elem(&tmp_pathinfos, &index, &path_info, BPF_ANY);
            return 0;
        } else {
            u64 index = is_src? MAX_PATH_DEPTH - 1 - depth : MAX_PATH_DEPTH * 2 - 1 - depth;
            bpf_map_update_elem(&tmp_pathinfos, &index, &path_info, BPF_ANY);
        }
        tmp_dentry = BPF_CORE_READ(tmp_dentry, d_parent);
        depth++;
    }
    return 0;
}

static int clear_path_info() {
    for (int i = 0; i < MAX_PATH_DEPTH * 2; i++) {
        u64 index = i;
        bpf_map_delete_elem(&tmp_pathinfos, &index);
    }
    return 0;
}

static int clear_mount_path_info() {
    bpf_printk("clear_mount_path_info");
    u64 index = 0;
    bpf_map_delete_elem(&tmp_mount_pathinfos, &index);
    return 0;
}

static int cache_mount_path_info(const char* name, u64 dev) {
    // 定义一个path_info结构体变量info
    struct path_info info;
    // 初始化info的深度为0
    info.depth = 0;
    // 初始化info的索引号为0
    info.ino = 0;
    // 初始化info的设备号为dev
    info.dev = dev;
    // 使用bpf_probe_read_kernel_str函数从内核空间读取name字符串到info.name中，读取的最大长度为NAME_MAX
    // 并返回读取的长度到info.len
    info.len = bpf_probe_read_kernel_str(info.name, NAME_MAX, name);
    // 定义一个u64类型的变量index并初始化为0
    u64 index = 0;
    // 使用bpf_printk函数打印挂载点的名称和设备号
    bpf_printk("cache mount name: %s, dev: %lu\n", info.name, dev);
    // 使用bpf_map_update_elem函数将index和info的映射关系更新到tmp_mount_pathinfos映射表中
    // BPF_ANY表示如果映射表中已有对应的key，则更新其value；否则插入新的key-value对
    bpf_map_update_elem(&tmp_mount_pathinfos, &index, &info, BPF_ANY);
    // 函数返回0，表示成功执行
    return 0;
}

/**
 * @brief 解析路径信息
 *
 * 从BPF映射中获取源路径和目标路径信息，并填充到事件结构体中。
 *
 * @param evt 事件结构体指针
 *
 * @return 返回值为0表示成功，其他值表示失败
 */
static int pharse_path_info(struct event *evt) {
    // 获取源路径
    // get src path
    u64 src_index = 0;
    u64 dest_index = MAX_PATH_DEPTH;
    struct path_info *src_root = bpf_map_lookup_elem(&tmp_pathinfos, &src_index);
    struct path_info *dest_root = bpf_map_lookup_elem(&tmp_pathinfos, &dest_index);

    if (src_root) {
        // 设置事件设备的设备号和源文件的inode号
        evt->dev = src_root->dev;
        evt->src_ino = src_root->ino;

        // 填充源路径信息
        for (int i = 0; i < MAX_PATH_DEPTH; i++) {
            u64 tmp = i;
            struct path_info *path_info = bpf_map_lookup_elem(&tmp_pathinfos, &tmp);
            if (path_info) {
                evt->src_path[i] = *path_info;
            }
        }
    }

    if (dest_root) {
        // 设置目标文件的inode号
        if (evt->type == FILE_RENAME) {
            struct dentry *tmp = dest_root->entry;
            if (!tmp) {
                evt->dst_ino = dest_root->ino;
            } else {
                evt->dst_ino = BPF_CORE_READ(tmp, d_inode, i_ino);
            }
        } else {
            evt->dst_ino = dest_root->ino;
        }
        // 填充目标设备号
        evt->dev = dest_root->dev;
        // 填充目标路径信息
        for (int i = MAX_PATH_DEPTH; i < MAX_PATH_DEPTH * 2; i++) {
            u64 tmp = i;
            struct path_info *path_info = bpf_map_lookup_elem(&tmp_pathinfos, &tmp);
            if (path_info) {
                evt->dst_path[i - MAX_PATH_DEPTH] = *path_info;
            }
        }
    }

    return 0;
}

static int handle_cache_and_submit(u32 type, int ret)
{
    // 在ring buffer中保留空间用于存储事件
    struct event *evt = bpf_ringbuf_reserve(&events, sizeof(struct event), 0);
    if (!evt) {
        // 如果预留空间失败，则打印错误信息并清除路径信息
        bpf_printk("handle_cache_and_submit: bpf_ringbuf_reserve failed %u\n", type);
        clear_path_info();
        return 0;
    }

    // 如果函数返回值不为0，表示创建文件失败
    if (ret != 0) {
        // 丢弃之前预留的空间
        bpf_ringbuf_discard(evt, 0);
        clear_path_info();
        return 0;
    }

    // 获取当前进程的PID和TGID
    u64 pid_tgid = bpf_get_current_pid_tgid();
    // 将TGID右移32位得到PID
    evt->pid = pid_tgid >> 32;
    // 获取当前进程的命令行名称
    bpf_get_current_comm(&evt->comm, sizeof(evt->comm));
    // 设置事件类型
    evt->type = type;
    // 解析路径信息
    pharse_path_info(evt);
    // 提交事件到ring buffer
    bpf_ringbuf_submit(evt, 0);
    // 清除路径信息
    clear_path_info();
    return 0;
}


SEC("kprobe/d_instantiate_new")
int BPF_KPROBE(vfs_create_entry, struct dentry *dentry, struct inode *inode) {
    // 调用cache_path_infos函数，将dentry和inode的i_ino值传递给函数
    int ret = cache_path_infos(dentry, BPF_CORE_READ(inode, i_ino), true);

    // 如果cache_path_infos函数返回非0值，表示操作失败
    if (ret != 0) {
        // 使用bpf_printk打印错误信息
        bpf_printk("vfs_create_entry: cache_path_infos failed\n");
    }

    // 函数返回0，表示操作成功
    return 0;
}

SEC("kretprobe/d_instantiate_new")
int BPF_KRETPROBE(vfs_create_exit) {
    // 调用handle_cache_and_submit函数，处理缓存和提交事件
    return handle_cache_and_submit(FILE_CREATE, 0);
}

SEC("kprobe/vfs_rename")
int BPF_KPROBE(vfs_rename_entry, struct renamedata *renamedata) {
    struct dentry *src_dentry = BPF_CORE_READ(renamedata, old_dentry);
    struct dentry *dst_dentry = BPF_CORE_READ(renamedata, new_dentry);
    u64 src_ino = BPF_CORE_READ(src_dentry, d_inode, i_ino);
    u64 dst_ino = BPF_CORE_READ(dst_dentry, d_inode, i_ino);

    // 调用cache_path_infos函数，将src dentry和inode的i_ino值传递给函数
    int ret = cache_path_infos(src_dentry, src_ino, true);
    // 如果cache_path_infos函数返回非0值，表示操作失败
    if (ret != 0) {
        // 使用bpf_printk打印错误信息
        bpf_printk("vfs_remove: cache_path_infos failed\n");
    }

    // 调用cache_path_infos函数，将dst dentry和inode的i_ino值传递给函数
    ret = cache_path_infos(dst_dentry, dst_ino, false);
    if (ret != 0) {
        bpf_printk("vfs_remove: cache_path_infos for dst_dentry failed\n");
    }

    // 函数返回0，表示操作成功
    return 0;
}

SEC("kretprobe/vfs_rename")
int BPF_KRETPROBE(vfs_rename_exit, int ret) {
    return handle_cache_and_submit(FILE_RENAME, ret);
}

SEC("kprobe/vfs_unlink")
int BPF_KPROBE(vfs_unlink_entry, struct mnt_idmap *idmap, struct inode *dir,
               struct dentry *dentry, struct inode **delegated_inode)
{
    // 调用cache_path_infos函数，将dentry和inode的i_ino值传递给函数
    int ret = cache_path_infos(dentry, BPF_CORE_READ(dir, i_ino), true);
    // 如果cache_path_infos函数返回非0值，表示操作失败
    if (ret != 0) {
        // 使用bpf_printk打印错误信息
        bpf_printk("vfs_unlink: cache_path_infos failed\n");
    }

    // 函数返回0，表示操作成功
    return 0;
}

SEC("kretprobe/vfs_unlink")
int BPF_KRETPROBE(vfs_unlink_exit, int ret) {
    // 调用handle_cache_and_submit函数，处理缓存和提交事件
    return handle_cache_and_submit(FILE_DELETE, ret);
}

SEC("kprobe/vfs_rmdir")
int BPF_KPROBE(vfs_rmdir_entry, struct mnt_idmap *idmap, struct inode *dir,
               struct dentry *dentry)
{
    // 调用cache_path_infos函数，将dentry和inode的i_ino值传递给函数
    int ret = cache_path_infos(dentry, BPF_CORE_READ(dir, i_ino), true);
    // 如果cache_path_infos函数返回非0值，表示操作失败
    if (ret != 0) {
        // 使用bpf_printk打印错误信息
        bpf_printk("vfs_rmdir: cache_path_infos failed\n");
    }

    // 函数返回0，表示操作成功
    return 0;
}

SEC("kretprobe/vfs_rmdir")
int BPF_KRETPROBE(vfs_rmdir_exit, int ret) {
    // 调用handle_cache_and_submit函数，处理缓存和提交事件
    return handle_cache_and_submit(FILE_DELETE, ret);
}

SEC("kprobe/file_modified")
int BPF_KPROBE(file_modified_entry, struct file *file) {
    struct dentry *dentry = BPF_CORE_READ(file, f_path.dentry);
    u64 ino = BPF_CORE_READ(dentry, d_inode, i_ino);

    int ret = cache_path_infos(dentry, ino, true);
    if (ret != 0) {
        bpf_printk("file_modified: cache_path_infos failed\n");
    }

    return 0;
}

SEC("kretprobe/file_modified")
int BPF_KRETPROBE(file_modified_exit, int ret) {
    return handle_cache_and_submit(FILE_CONTENT_CHANGED, ret);
}

SEC("kprobe/notify_change")
int BPF_KPROBE(notify_change_entry, struct mnt_idmap *idmap, struct dentry *dentry,
               struct iattr *attr, struct inode **delegated_inode)
{
    // 读取ia_valid字段的值
    unsigned int ia_valid = BPF_CORE_READ(attr, ia_valid);

    // 初始化索引和权限变化标志
    u64 index = 0;
    u64 is_permission_changed = 0;

    // 检查ia_valid是否包含ATTR_MODE、ATTR_UID或ATTR_GID
    if ((ia_valid & (ATTR_MODE | ATTR_UID | ATTR_GID))) {
        unsigned short iamode = BPF_CORE_READ(attr, ia_mode);
        unsigned short currentmode = BPF_CORE_READ(dentry, d_inode, i_mode);
        if (iamode == currentmode) {
            // 如果ia_mode和当前模式相同，则返回
            bpf_map_update_elem(&permission_cache, &index, &is_permission_changed, BPF_ANY);
            return 0;
        }

        kuid_t iauid = BPF_CORE_READ(attr, ia_uid);
        kuid_t currentuid = BPF_CORE_READ(dentry, d_inode, i_uid);
        if (iauid.val == currentuid.val) {
            // 如果ia_uid和当前用户ID相同，则返回
            bpf_map_update_elem(&permission_cache, &index, &is_permission_changed, BPF_ANY);
            return 0;
        }

        kgid_t iagid = BPF_CORE_READ(attr, ia_gid);
        kgid_t currentgid = BPF_CORE_READ(dentry, d_inode, i_gid);
        if (iagid.val == currentgid.val) {
            // 如果ia_gid和当前组ID相同，则返回
            bpf_map_update_elem(&permission_cache, &index, &is_permission_changed, BPF_ANY);
            return 0;
        }

        // 权限改变
        is_permission_changed = 1;
        // 更新权限缓存
        bpf_map_update_elem(&permission_cache, &index, &is_permission_changed, BPF_ANY);

        // 获取inode编号
        u64 ino = BPF_CORE_READ(dentry, d_inode, i_ino);
        // 缓存路径信息
        int ret = cache_path_infos(dentry, ino, true);
        if (ret != 0) {
            // 如果缓存路径信息失败，则打印错误信息
            bpf_printk("notify_change: cache_path_infos failed\n");
        }
        return 0;
    }

    // 更新权限缓存
    bpf_map_update_elem(&permission_cache, &index, &is_permission_changed, BPF_ANY);

    return 0;
}

SEC("kretprobe/notify_change")
int BPF_KRETPROBE(notify_change_exit, int ret) {
    u64 index = 0;
    // 在permission_cache中查找index对应的元素
    u64 *is_permission_changed = bpf_map_lookup_elem(&permission_cache, &index);
    if (is_permission_changed == NULL) {
        // 如果查找结果为空，打印错误信息
        bpf_printk("notify_change_exit: is_permission_changed is NULL\n");
        return 0;
    }
    if (*is_permission_changed == 0) {
        return 0;
    }
    // 处理缓存并提交事件
    handle_cache_and_submit(FILE_PERM_CHANGE, ret);
    // 清除路径信息
    clear_path_info();
    return 0;
}

int handle_mount_event_and_submit(u32 type, int ret)
{
    if (ret) {
        clear_mount_path_info();
        return 0;
    }

    struct event *evt = bpf_ringbuf_reserve(&events, sizeof(struct event), 0);
    if (!evt) {
        bpf_printk("handle_mount_event_and_submit: bpf_ringbuf_reserve failed\n");
        return 0;
    }

    u64 index = 0;
    struct path_info *path = bpf_map_lookup_elem(&tmp_mount_pathinfos, &index);
    if (!path) {
        bpf_printk("handle_mount_event_and_submit: bpf_map_lookup_elem failed\n");
        bpf_ringbuf_discard(evt, 0);
        return 0;
    }
    // 获取当前进程的PID和TGID
    u64 pid_tgid = bpf_get_current_pid_tgid();
    // 将TGID右移32位得到PID
    evt->pid = pid_tgid >> 32;
    // 获取当前进程的命令行名称
    bpf_get_current_comm(&evt->comm, sizeof(evt->comm));
    // 设置事件类型
    evt->type = type;
    evt->dev = path->dev;
    evt->src_path[0] = *path;
    bpf_ringbuf_submit(evt, 0);
    bpf_printk("handle_mount_event_and_submit %s %lu", path->name, path->dev);
    return 0;
}

SEC("kprobe/do_new_mount")
int BPF_KPROBE(do_new_mount_entry, struct path *path, const char *fstype, int sb_flags,
               int mnt_flags, const char *name, void *data) {
    // fixme: 使用bpf map在kretprobe中传递ringbuf，目前存在高概率数据异常问题，暂时使用直接在entry中发送事件的方式
    struct event *evt = bpf_ringbuf_reserve(&events, sizeof(struct event), 0);
    if (!evt) {
        bpf_printk("do_new_mount: bpf_ringbuf_reserve failed\n");
        return 0;
    }

    evt->pid = bpf_get_current_pid_tgid() >> 32;
    bpf_get_current_comm(&evt->comm, sizeof(evt->comm));
    evt->type = MOUNT_ADDED;
    evt->dev = BPF_CORE_READ(path, mnt, mnt_sb, s_dev);
    evt->src_path[0].len = bpf_probe_read_str(&evt->src_path[0].name, sizeof(evt->src_path[0].name), name);
    
    bpf_ringbuf_submit(evt, 0);

    return 0;
}

SEC("kretprobe/do_new_mount")
int BPF_KRETPROBE(do_new_mount_exit, int ret) {
    return 0;
    //return handle_mount_event_and_submit(MOUNT_ADDED, ret);
}

SEC("kprobe/do_umount")
int BPF_KPROBE(do_umount_entry, struct mount *mnt, int flags) {
    clear_mount_path_info();
    cache_mount_path_info(BPF_CORE_READ(mnt, mnt_devname), BPF_CORE_READ(mnt, mnt.mnt_sb, s_dev));
    return 0;
}

SEC("kretprobe/do_umount")
int BPF_KRETPROBE(do_umount_exit, int ret) {
    return handle_mount_event_and_submit(MOUNT_REMOVED, ret);
}
